package scu.maqiang.numeric;

import java.util.Arrays;

/**
 * 行压缩稀疏矩阵类
 * @author 马强
 *
 */
public class CSRCMatrix {
	//稀疏矩阵行指标
	private int[] rowIndex;
	//稀疏矩阵列指标
	private int[] columns;
	//稀疏矩阵元素值
	private Complex[] values;
	//稀疏矩阵行数
	private int m;
	//稀疏矩阵列数
	private int n;
	//稀疏矩阵非0元素
	private int nnz;
	//稀疏矩阵是否对称
	private boolean isHermitian;
	
	/**
	 * 得到矩阵行数
	 * @return 矩阵行数
	 */
	public int getM() {
		return m;
	}

	/**
	 * 得到矩阵列数
	 * @return 矩阵列数
	 */
	public int getN() {
		return n;
	}

	/**
	 * 得到矩阵非零元素个数
	 * @return 矩阵非零元素个数
	 */
	public int getNNZ() {
		return nnz;
	}

	/**
	 * 得到矩阵是否对称
	 * @return 矩阵是否对称
	 */
	public boolean isHermitian() {
		return isHermitian;
	}

	/**
	 * 设置矩阵对称性，可将上三角矩阵设置为对称属性
	 * @param isHermitian 矩阵对称性
	 */
	public void setHermitian(boolean isHermitian) {
		this.isHermitian = isHermitian;
	}
	
	/**
	 * 构造函数，给出矩阵行列数构造矩阵
	 * @param m 矩阵行数
	 * @param n 矩阵列数
	 */
	public CSRCMatrix(int m, int n) {
		this.m = m;
		this.n = n;
		this.nnz = 0;
		rowIndex = null;
		columns = null;
		values = null;
		isHermitian = false;
	}
	
	/**
	 * 构造函数，给出方阵阶数构造矩阵
	 * @param n 矩阵阶数
	 */
	public CSRCMatrix(int n) {
		this(n, n);
	}
	
	/**
	 * 构造函数，由SCMatrix类构造矩阵
	 * @param mat SRMmatrix对象
	 */
	public CSRCMatrix(SCMatrix mat) {
		if (mat.isSorted())
			mat.sort();
		isHermitian = mat.isHermitian();
		m = mat.getM();
		n = mat.getN();
		rowIndex = new int[m + 1];
		int colIdx = 0;
		int rowIdx;
		if (isHermitian) {
			nnz = (mat.getNNZ() + mat.getM()) / 2;
			columns = new int[nnz];
			values = new Complex[nnz];
			for (int i = 0; i < m; i++) {
				rowIdx = 0;
				for (SCEntry ele : mat.mat[i]) {
					if (ele.col >= i) {
						columns[colIdx] = ele.col;
						values[colIdx].real = ele.val.real;
						values[colIdx].image = ele.val.image;
						colIdx++;
						rowIdx++;
					}
				}
				rowIndex[i + 1] = rowIndex[i] + rowIdx;
			}
		} else {
			nnz = mat.getNNZ();
			columns = new int[nnz];
			values = new Complex[nnz];
			
			for (int i = 0; i < m; i++) {
				for (SCEntry ele : mat.mat[i]) {
					columns[colIdx] = ele.col;
					values[colIdx].real = ele.val.real;
					values[colIdx].image = ele.val.image;
					colIdx++;
				}
				rowIndex[i + 1] = rowIndex[i] + mat.mat[i].size();
			}
		}
	}
	
	
	/**
	 * 得到矩阵行列指标下的元素值
	 * @param i 行指标
	 * @param j 列指标
	 * @return 元素值
	 */
	public Complex getElement(int i, int j) {
		ParamCheck.checkBounds(i, 0, m);
		ParamCheck.checkBounds(j, 0, n);
		for(int k = rowIndex[i], kEnd = rowIndex[i + 1]; k < kEnd; k++) {
			if(columns[k] == j) {
				return values[k];
			}
		}
		return new Complex(0, 0);
	}
	
	/**
	 * 将SCMatrix对象转换为CSRCMatrix对象, 对任意一个SCMatrix矩阵A，该静态方法将其转换为两个CSRCMatrix: mat1, mat2, 
	 * 以及用一个数组diag得到该矩阵的对角元素, 如果矩阵A对称，则矩阵mat1为该对称矩阵的上三角元素，mat2为矩阵下三角元素但不包含主对角元素;
	 * 如果矩阵B不对称，则得到mat1为整个矩阵, mat2仍为原矩阵下三角元素但不包含主对角线元素构成的矩阵
	 * @param mat 一般SCMatrix矩阵
	 * @param diag 得到矩阵对角线元素
	 * @param mat1 CSRCMatrix矩阵，当mat对称时，用mat上三角元素形成矩阵，并设置其为对称的; 当mat不对称时，用mat所有元素形成矩阵
	 * @param mat2 CSRCMatrix矩阵，用mat下三角元素形成矩阵
	 */
	public static void SCMatrixToCSRCMatrix(SCMatrix mat, Complex[] diag, CSRCMatrix mat1, CSRCMatrix mat2) {
		boolean isSymmetric = mat.isHermitian();
		int colIdxL = 0;
		int rowIdxL = 0;
		if(isSymmetric) {
			extractDLU(mat, diag, mat1, mat2);
		}else {
			if (!mat.isSorted())
				mat.sort();
			mat1.isHermitian = isSymmetric;
			int m = mat.getM();
			int n = mat.getN();
			mat1.m = m;
			mat1.n = n;
			mat1.rowIndex = new int[m + 1];
			int colIdx = 0;
			int nnz = mat.getNNZ();
			mat1.nnz = nnz;
			mat1.columns = new int[nnz];
			mat1.values = Complex.zeros(nnz);
			
			
			int NNZL = mat.getNNZL();
			mat2.m = m;
			mat2.n = n;
			mat2.nnz = NNZL;
			mat2.rowIndex = new int[m + 1];
			mat2.columns = new int[NNZL];
			mat2.values = Complex.zeros(NNZL);
			for (int i = 0; i < m; i++) {
				rowIdxL = 0;
				for (SCEntry ele : mat.mat[i]) {
					int colj = ele.col;
					mat1.columns[colIdx] = colj;
					mat1.values[colIdx].real = ele.val.real;
					mat1.values[colIdx].image = ele.val.image;
					colIdx++;
					if(colj < i) { // lower entry
						mat2.columns[colIdxL] = colj;
						mat2.values[colIdxL].real = ele.val.real;
						mat2.values[colIdxL].image = ele.val.image;
						rowIdxL++;
						colIdxL++;
					}
					if (colj == i) {
						diag[i].real = ele.val.real;
						diag[i].image = ele.val.image;
					}
				}
				mat1.rowIndex[i + 1] = mat1.rowIndex[i] + mat.mat[i].size();
				mat2.rowIndex[i + 1] = mat2.rowIndex[i] + rowIdxL;
			}
		}
	}

	/**
	 * 将SCMatrix对象转换为CSRCMatrix对象, 对任意一个SCMatrix矩阵A，该静态方法将其转换为两个CSRCMatrix: mat1, mat2,
	 * 以及用一个数组diag得到该矩阵的对角元素, 如果矩阵A对称，则矩阵mat1为该对称矩阵的上三角元素，mat2为矩阵下三角元素但不包含主对角元素;
	 * 如果矩阵B不对称，则得到mat1为整个矩阵, mat2仍为原矩阵下三角元素但不包含主对角线元素构成的矩阵
	 * @param mat 一般SCMatrix矩阵
	 * @param diag 得到矩阵对角线元素
	 * @param mat1 CSRCMatrix矩阵，当mat对称时，用mat上三角元素形成矩阵，并设置其为对称的; 当mat不对称时，用mat所有元素形成矩阵
	 * @param mat2 CSRCMatrix矩阵，用mat下三角元素形成矩阵
	 */
	public static void SCMatrixToCSRCMatrix(SCMatrix mat, Complex[] diag, CSRCMatrix mat1, CSRCMatrix mat2, boolean isUpdate) {
		if (isUpdate) {
			SCMatrixToCSRCMatrix(mat, diag, mat1, mat2);
			return;
		}
		boolean isSymmetric = mat.isHermitian();
		int colIdxL = 0;
		int rowIdxL = 0;
		if(isSymmetric) {
			extractDLU(mat, diag, mat1, mat2, isUpdate);
		}else {
//			if (!mat.isSorted())
//				mat.sort();
//			mat1.isHermitian = isSymmetric;
			int m = mat.getM();
			int n = mat.getN();
//			mat1.m = m;
//			mat1.n = n;
//			mat1.rowIndex = new int[m + 1];
			int colIdx = 0;
//			int nnz = mat.getNNZ();
//			mat1.nnz = nnz;
//			mat1.columns = new int[nnz];
//			mat1.values = Complex.zeros(nnz);
//
//			int NNZL = mat.getNNZL();
//			mat2.m = m;
//			mat2.n = n;
//			mat2.nnz = NNZL;
//			mat2.rowIndex = new int[m + 1];
//			mat2.columns = new int[NNZL];
//			mat2.values = Complex.zeros(NNZL);
			for (int i = 0; i < m; i++) {
				rowIdxL = 0;
				for (SCEntry ele : mat.mat[i]) {
					int colj = ele.col;
					mat1.columns[colIdx] = colj;
					mat1.values[colIdx].real = ele.val.real;
					mat1.values[colIdx].image = ele.val.image;
					colIdx++;
					if(colj < i) { // lower entry
						mat2.columns[colIdxL] = colj;
						mat2.values[colIdxL].real = ele.val.real;
						mat2.values[colIdxL].image = ele.val.image;
						rowIdxL++;
						colIdxL++;
					}
					if (colj == i) {
						diag[i].real = ele.val.real;
						diag[i].image = ele.val.image;
					}
				}
				mat1.rowIndex[i + 1] = mat1.rowIndex[i] + mat.mat[i].size();
				mat2.rowIndex[i + 1] = mat2.rowIndex[i] + rowIdxL;
			}
		}
	}
	
	/**
	 * 对于一个SCMatrix对象的对称矩阵，将其转换为两个CSRCMatrix，并得到其对角元素
	 * @param mat 对称SCMatrix矩阵
	 * @param diag 得到矩阵对角线元素
	 * @param matU 得到矩阵上三角元素，并设置矩阵对称
	 * @param matL 矩阵矩阵下三角元素，但不包含主对角线元素
	 */
	public static void extractDLU(SCMatrix mat, Complex[] diag, CSRCMatrix matU, CSRCMatrix matL) {
		int m = mat.getM();
		ParamCheck.checkEqual(diag.length, m);
		ParamCheck.checkEqual(matU.m, m);
		ParamCheck.checkEqual(matL.m, m);
		if (!mat.isSorted()) {
			mat.sort();
		}
		int NNZL = mat.getNNZL();
		int NNZU = mat.getNNZU();
		int rowIdxL, rowIdxU;
		int colIdxL = 0, colIdxU = 0;
		matU.nnz = NNZU + m;
		matU.rowIndex = new int[m + 1];
		matU.columns = new int[NNZU + m];
		matU.values = Complex.zeros(NNZU + m);
		
		matL.nnz = NNZL;
		matL.rowIndex = new int[m + 1];
		matL.columns = new int[NNZL];
		matL.values = Complex.zeros(NNZL);
		for(int i = 0; i < m; i++) {
			rowIdxL = 0;
			rowIdxU = 0;
			for (SCEntry ele: mat.mat[i]) {
				int colj = ele.col;
				if(colj < i) { // lower entry
					matL.columns[colIdxL] = colj;
					matL.values[colIdxL].real = ele.val.real;
					matL.values[colIdxL].image = ele.val.image;
					rowIdxL++;
					colIdxL++;
				}else { // upper plus diag entry
					matU.columns[colIdxU] = colj;
					matU.values[colIdxU].real = ele.val.real;
					matU.values[colIdxU].image = ele.val.image;
					rowIdxU++;
					colIdxU++;					
				}
				if (colj == i) { //diag entry
					diag[i].real = ele.val.real;
					diag[i].image = ele.val.image;
				}
			}
			matL.rowIndex[i + 1] = matL.rowIndex[i] + rowIdxL;
			matU.rowIndex[i + 1] = matU.rowIndex[i] + rowIdxU;
		}
		matU.setHermitian(true);
	}

	public static void extractDLU(SCMatrix mat, Complex[] diag, CSRCMatrix matU, CSRCMatrix matL, boolean isUpdate) {
		if(isUpdate) {
			extractDLU(mat, diag, matU, matL);
			return;
		}
		if (!mat.isSorted()) {
			mat.sort();
		}
		int m = mat.getM();
		int rowIdxL, rowIdxU;
		int colIdxL = 0, colIdxU = 0;
		for(int i = 0; i < m; i++) {
			rowIdxL = 0;
			rowIdxU = 0;
			for (SCEntry ele: mat.mat[i]) {
				int colj = ele.col;
				if(colj < i) { // lower entry
					matL.columns[colIdxL] = colj;
					matL.values[colIdxL].real = ele.val.real;
					matL.values[colIdxL].image = ele.val.image;
					rowIdxL++;
					colIdxL++;
				}else { // upper plus diag entry
					matU.columns[colIdxU] = colj;
					matU.values[colIdxU].real = ele.val.real;
					matU.values[colIdxU].image = ele.val.image;
					rowIdxU++;
					colIdxU++;
				}
				if (colj == i) { //diag entry
					diag[i].real = ele.val.real;
					diag[i].image = ele.val.image;
				}
			}
			matL.rowIndex[i + 1] = matL.rowIndex[i] + rowIdxL;
			matU.rowIndex[i + 1] = matU.rowIndex[i] + rowIdxU;
		}
		matU.setHermitian(true);
	}
	

	/**
	 * 矩阵与向量乘积运算
	 * @param u 向量 
	 * @return 结果向量
	 */
	public Complex[] mul(Complex[] u) {
		Complex[] v = new Complex[m];
		for(int i = 0; i < m; i++) {
			v[i] = new Complex(0.0, 0.0);
		}
		mmul(u, v);
		return v;
	}
	
	/**
	 * 矩阵向量乘积运算
	 * @param u 向量
	 * @param v 结果向量
	 */
	public void mul(Complex[] u, Complex[] v) {
		ParamCheck.checkEqual(v.length, m);
		for(int i = 0; i < m; i++) {
			v[i].real = 0.0;
			v[i].image = 0.0;
		}
		mmul(u, v);
		//System.out.println(Arrays.toString(v));
	}
	
	/**
	 * 
	 * @param u
	 * @param v
	 */
	private void mmul(Complex[] u, Complex[] v) {
		int i, j, jEnd, colj;
		Complex tt, t2;
		if (isHermitian) {
			for (i = 0; i < m; i++) {
				jEnd = rowIndex[i + 1];
				for(j =  rowIndex[i]; j < jEnd; j++) {
					colj = columns[j];
					tt = values[j].mul(u[colj]);
					v[i].real += tt.real;
					v[i].image += tt.image;
					if (i != colj) {
						t2 = u[i].mulT(values[j]);
						v[colj].real += t2.real;
						v[colj].image += t2.image;
					}
				}
			}
		}else {
			for (i = 0; i < m; i++) {
				jEnd = rowIndex[i + 1];
				for (j = rowIndex[i]; j < jEnd; j++) {
					tt = values[j].mul(u[columns[j]]);
					v[i].real += tt.real;
					v[i].image += tt.image;
				}
			}
		}
	}

	/**
	 * 矩阵转置与向量乘积
	 * @param u 向量
	 * @return 结果向量
	 */
	public Complex[] mulT(Complex[] u) {
		Complex[] v = new Complex[m];
		for(int i = 0, n = v.length; i < n; i++) {
			v[i] = new Complex(0.0, 0.0);
		}
		mmulT(u, v);
		return v;
	}
	
	/**
	 * 
	 * @param u
	 * @param v
	 */
	public void mulT(Complex[] u, Complex[] v) {
		ParamCheck.checkEqual(v.length, m);
		for(int i = 0; i < m; i++) {
			v[i].real = 0.0;
			v[i].image = 0.0;
		}
		mmulT(u, v);
	}
	
	private void mmulT(Complex[] u, Complex[] v) {
		int i, j, jEnd, colj;
		if (isHermitian) {
			for (i = 0; i < m; i++) {
				jEnd = rowIndex[i + 1];
				for(j =  rowIndex[i]; j < jEnd; j++) {
					colj = columns[j];
					Complex tt = u[colj].mulT(values[j]);
					v[i].real += tt.real;
					v[i].image += tt.image;
					if (i != colj) {
						Complex t2 = u[i].mul(values[j]);
						v[colj].real += t2.real;
						v[colj].image += t2.image;
					}
				}
			}
		}else {
			for (i = 0; i < m; i++) {
				jEnd = rowIndex[i + 1];
				for (j = rowIndex[i]; j < jEnd; j++) {
					Complex t2 = values[j].mulT(u[i]);
					v[columns[j]].real += t2.real;
					v[columns[j]].image += t2.image;
				}
			}
		}		
	}
	
	/**
	 * 
	 * @param diag
	 * @param matU
	 * @param matL
	 * @param r0
	 * @param z0
	 * @param omega
	 */
	public static void SSORTreat(Complex[] diag, CSRCMatrix matU, CSRCMatrix matL, Complex[] r0, Complex[] z0, double omega) {
		int n = matL.getN();
		int jEnd;
		Complex temp = new Complex(0.0, 0.0);
		r0[0].div(diag[0], z0[0]);
		Complex tt = new Complex(0.0, 0.0);
		for (int i = 1; i < n; i++) {
			temp.real = 0.0;
			temp.image = 0.0;
			jEnd = matL.rowIndex[i + 1];
			for(int j = matL.rowIndex[i]; j < jEnd; j++) {
				matL.values[j].mul(z0[matL.columns[j]], tt);
				temp.real += omega * tt.real;
				temp.image += omega * tt.image;
			}
			//System.out.println(r0[i] + "\t" + temp + "\t" + diag[i]);
			r0[i].sub(temp, tt);
			tt.div(diag[i], z0[i]);
			//z0[i] = r0[i].sub(temp).div(diag[i]);
			//if(Double.isNaN(z0[i])) {
			//	System.exit(1);
			//}
		}
		
		//System.out.println("z0 = " + MVO.toString(z0));
		for(int i = 0; i < n; i++) {
			z0[i].mul(diag[i], tt);
			z0[i].real = tt.real;
			z0[i].image = tt.image;
		}
		
		z0[n - 1].div(diag[n - 1], tt);
		z0[n - 1].real = tt.real;
		z0[n - 1].image = tt.image;
		for (int i = n - 2; i >= 0; i--) {
			temp.real = 0.0;
			temp.image = 0.0;
			jEnd = matU.rowIndex[i + 1];
			for (int j = matU.rowIndex[i]; j < jEnd; j++) {
				if (matU.columns[j] > i) {
					matU.values[j].mul(z0[matU.columns[j]], tt);
					temp.real += omega * tt.real;
					temp.image += omega * tt.image;
				}
			}
			z0[i].sub(temp, tt);
			tt.div(diag[i], z0[i]);
//			z0[i] = z0[i].sub(temp).div(diag[i]);
		}
}

	
//	public static void GSPrecondition(double[] diag, CSRCMatrix matU, double[] r0, double[] z0) {
//		int n = matU.getN();
//		int jEnd;
//		double temp;
//		z0[0] = r0[0] / diag[0];
//		for (int i = 1; i < n; i++) {
//			temp = 0.0;
//			jEnd = matU.rowIndex[i + 1];
//			for(int j = matU.rowIndex[i]; j < jEnd; j++) {
//				if (matU.columns[j] != i) {
//					temp += matU.values[j] * z0[matU.columns[j]]; 
//				}
//			}
//			z0[i] = (r0[i] - temp) / diag[i];
//		}
//	}
	
	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder();
		sb.append(m + "\t" + n + "\t" + nnz + "\t" + isHermitian +"\n");
		sb.append(Arrays.toString(rowIndex) + "\n");
		sb.append(Arrays.toString(columns) + "\n");
		sb.append(Arrays.toString(values) + "\n");
		return sb.toString();
	}
	
	/**
	 * 
	 * @return
	 */
	public String toSRFormat() {
		StringBuilder sb = new StringBuilder();
		sb.append(m + "\t" + n + "\t" + nnz + "\t" + isHermitian +"\n");
		int jEnd, idx;
		idx = 0;
		for(int i = 0; i < m; i++) {
			jEnd = rowIndex[i + 1];
			for(int j = rowIndex[i]; j < jEnd; j++) {
				sb.append(String.format("%1$6d,%2$6d,%3$16.8f\n", i, columns[idx], values[idx]));
				idx++;
			}
		}
		return sb.toString();
	}

	public void clear() {
		rowIndex = null;
		columns = null;
		if (values != null) {
			for(int i = 0, n = values.length; i < n; i++) {
				values[i] = null;
			}
		}
		values = null;
		m = 0;
		n = 0;
		nnz = 0;
		isHermitian = false;
	}
}
